import java.net.URL;

import org.openqa.selenium.Capabilities;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.WebDriverException;
import org.openqa.selenium.remote.CapabilityType;
import org.openqa.selenium.remote.DriverCommand;
import org.openqa.selenium.remote.RemoteWebDriver;

public class RemoteScreenshot extends RemoteWebDriver {

	public RemoteScreenshot(URL remoteAddress, Capabilities desiredCapabilities) {
		super(remoteAddress, desiredCapabilities);
	}

	public <X> X getScreenshotAs(OutputType<X> target)
			throws WebDriverException {
		boolean criteria = suopportsScreenshot();
		if (criteria) {
			String remoteResult = execute(DriverCommand.SCREENSHOT).getValue()
					.toString();
			return target.convertFromBase64Png(remoteResult);
		}
		return null;
	}

	private boolean suopportsScreenshot() {
		boolean criteria = (Boolean) getCapabilities().getCapability(
				CapabilityType.TAKES_SCREENSHOT);
		return criteria;
	}
}
